#include "objs/sinx_ispc.h"
#include <iostream>
#include <memory>
#include <utility>
#include <chrono>
#include <immintrin.h>
#include <cstring>
#include <cmath>

float sum_summary_avx(int N, float x[])
{
    float temp[8];
    __m256 partial = _mm256_setzero_ps();
    
    int i;
    for(i = 0; i <= N - 8; i += 8)
        partial = _mm256_add_ps(partial, _mm256_loadu_ps(&x[i]));
    
    // Handle remaining elements if N is not a multiple of 8
    float sum = 0.0f;
    for (; i < N; i++)
        sum += x[i];
    
    _mm256_storeu_ps(temp, partial);
    for(int j = 0; j < 8; j++)
        sum += temp[j];
    
    return sum;
}

float sum_matrix_avx(int width, int height, float matrix[][257]) 
{ 
    float sum = 0.0f;
    for (int j = 0; j < height; j++) {
        sum += sum_summary_avx(width, matrix[j]);
    }
    return sum;
}

void test_ispc_iteration()
{
    constexpr int n = 512;
    constexpr int N = n + 2;
    float* x = new float[N * N];
    double tolerance = 0.1;

    bool done = false;
    float diff = 0.0f;
    bool red = false;
    while(!done)
    {
        diff = ispc::iteration_grid(N, n, x, red);
        red = !red;
        if(diff / (n * n) < tolerance)
            done = true;
    }
    delete[] x;
}

void test_sequential_iteration()
{
    constexpr int n = 512;
    constexpr int N = n + 2;
    float* x = new float[N * N];
    double tolerance = 0.1;

    bool done = false;
    float diff = 0.0f;
    bool red = false;
    
    int row, col;
    float old, up, down, left, right, update;
    while(!done)
    {
        diff = 0.0f;
        for(int i = 0; i < N; i ++)
        {
            for(int j = 0; j < N; j ++)
            {
                if(i == 0 || j == 0 || i == N - 1 || j == N - 1) continue;
                old = x[i * N + j];
                up = x[(i - 1) * N + j];
                down = x[(i + 1) * N + j];
                left = x[i * N + j - 1];
                right = x[i * N + j + 1];
                update = (up + down + left + right) * 0.2f;
                if(red)
                {
                    if(i % 2 == 1 && j % 2 == 1) x[i * N + j] = update;
                    if(i % 2 == 0 && j % 2 == 0) x[i * N + j] = update;
                }
                else
                {
                    if(i % 2 == 1 && j % 2 == 0) x[i * N + j] = update;
                    if(i % 2 == 0 && j % 2 == 1) x[i * N + j] = update;
                }
                diff += std::fabs(x[i * N + j] - old);
            }
        }
        
        red = !red;
        if(diff / (n * n) < tolerance)
            done = true;
    }
    delete[] x;
}

int main()
{
    int N = 1024 * 4;
    int terms = 5;
    float* x = new float[N];
    float* result = new float[N];

    for(int i = 0; i < N; i ++) x[i] = 2;

    auto ispc_start = std::chrono::steady_clock::now();
    ispc::ispc_sinx_v3(N, terms, x, result);
    auto ispc_end = std::chrono::steady_clock::now();
    std::cout << "ispc time: " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;

    std::cout << "ispc reduce" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    float ret = ispc::sum_array(N, x);
    ispc_end =   std::chrono::steady_clock::now();
    std::cout << "ispc reduce time: " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;
    std::cout << "ispc res is: " << ret << std::endl;
    
    // for(int i = 0; i < N; i ++) x[i] = 2;
    std::cout << "AVX reduce" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    float ret2 = sum_summary_avx(N, x);
    ispc_end =   std::chrono::steady_clock::now();
    std::cout << "AVX reduce time: " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;
    std::cout << "AVX res is: " << ret2 << std::endl;

    int row = 256;
    int col = 256;
    float matrix[257][257];
    for (int i = 0; i < row; i++) 
    {
        for (int j = 0; j < col; j++) {
            matrix[i][j] = 1.0f;
        }
    }
    std::cout << "AVX Matrix sum" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    float ret3 = sum_matrix_avx(col, row, matrix);
    ispc_end   =   std::chrono::steady_clock::now();
    std::cout << "AVX Matrix sum time : " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;
    std::cout << "AVX Matrix res is: " << ret3 << std::endl;

    std::cout << "ISPC Matrix sum" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    float ret4 = ispc::sum_matrix(col, row, matrix);
    ispc_end   =   std::chrono::steady_clock::now();
    std::cout << "ISPC Matrix sum time : " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;
    std::cout << "ISPC Matrix res is: " << ret4 << std::endl;

    std::cout << "Matrix Sum Sequential" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    float ret5 = 0;
    for (int i = 0; i < row; i++)
    {
        for (int j = 0; j < col; j++)
        {
            ret5 += matrix[i][j];
        }
    }
    ispc_end   =   std::chrono::steady_clock::now();
    std::cout << "Matrix Sum Sequential time : " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;
    std::cout << "Matrix Sum Sequential res is: " << ret5 << std::endl;

    std::cout << "test ispc iteration" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    test_ispc_iteration();
    ispc_end   =   std::chrono::steady_clock::now();
    std::cout << "test ispc iteration time : " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;
    
    std::cout << "test sequential iteration" << std::endl;
    ispc_start = std::chrono::steady_clock::now();
    test_sequential_iteration();
    ispc_end   =   std::chrono::steady_clock::now();
    std::cout << "test sequential iteration time : " << std::chrono::duration_cast<std::chrono::nanoseconds>(ispc_end - ispc_start).count() << "ns" << std::endl;

    
    return 0;
}